import random

import gym
import numpy as np
import sounddevice as sd
from gym import spaces
from stable_baselines3.ddpg.ddpg import DDPG

from utils import util
from utils.simpledaw import SimpleDAW


class Simulator(gym.Env):
    def __init__(self):
        self.action_space = spaces.Box(high=1., low=0., shape=(244,))
        self.observation_space = spaces.Box(high=1., low=0., shape=(244,))
        self.daw = SimpleDAW(plugin='C:/VstPlugins/64bit/Sylenth1.dll', sample_rate=44100, bpm=120)
        self.midi_paths = util.get_files(r'B:\muse_repo\MIDI', 'mid')
        self.step_count = 0

    def step(self, action):
        self.daw.load_midi(random.choice(self.midi_paths))
        self.daw.set_params(list(action))
        print('正在播放预测声音')
        sd.play(self.daw.render(4.), 44100, blocking=True)
        print('播放结束')
        
        state = np.array(list(action))
        # print(state.shape)

        '''if action == 2:
            reward = 1
        else:
            reward = -1'''
        
        reward = float(input('请输入奖励：'))
        
        self.step_count += 1
        print('[第{}轮训练]'.format(self.step_count), 'action:', action[:10], 'state:', state[:10])
        
        done = True
        info = {}
        return state, reward, done, info

    def reset(self):
        state = np.zeros(244)
        return state

    def render(self, mode='human'):
        pass

    def seed(self, seed=None):
        pass


if __name__ == "__main__":
    env = Simulator()

    model = DDPG(policy="MlpPolicy", env=env)
    for _ in range(100):
        try:
            model.learn(total_timesteps=1)
            model.save('pg_rl.zip')
        except Exception:
            continue

    obs = env.reset()
    # 验证十次
    for _ in range(10):
        action, state = model.predict(observation=obs)
        obs = env.reset()
        print(action[:10])
        obs, reward, done, info = env.step(action)
        env.render()
